import requests
import json
from loguru import logger

# The URL for the ui-tars-server
SERVER_URL = "http://127.0.0.1:22001/process_task"

class UIAgent:
    def __init__(self):
        logger.info("UIAgent initialized.")

    def send_task_request(self, prompt, max_step=100, max_image=1):
        """
        Sends a task request to the ui-tars-server.
        """
        params = {
            "prompt": prompt,
            "max_step": max_step,
            "max_image": max_image
        }

        logger.info(f"Sending GET request to {SERVER_URL} with params: {params}")
        try:
            response = requests.get(SERVER_URL, params=params)
            response.raise_for_status()  # Raise an HTTPError for bad responses (4xx or 5xx)
            return response.json()
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Connection Error: Could not connect to the server at {SERVER_URL}. Is it running? {e}")
            raise Exception(f"Could not connect to UI-TARS server: {e}")
        except requests.exceptions.HTTPError as e:
            logger.error(f"HTTP Error: {e.response.status_code} - {e.response.text}")
            raise Exception(f"HTTP Error from UI-TARS server: {e.response.status_code} - {e.response.text}")
        except Exception as e:
            logger.error(f"An unexpected error occurred during UI-TARS request: {e}")
            raise Exception(f"An unexpected error occurred: {e}")

    def process_command(self, request_json: str) -> str:
        """
        Processes a JSON request to interact with the UI-TARS server.
        """
        logger.info(f"UIAgent received request: {request_json}")

        try:
            request = json.loads(request_json)
            task = request.get("task", {})
            prompt = task.get("prompt")
            max_step = task.get("max_step", 100)
            max_image = task.get("max_image", 1)

            if not prompt:
                return self._format_error_response(request.get("task_id"), "Missing 'prompt' in request task.")

            result = self.send_task_request(prompt, max_step, max_image)
            return self._format_success_response(request.get("task_id"), result)

        except json.JSONDecodeError:
            return self._format_error_response(None, "Invalid JSON request.")
        except Exception as e:
            logger.error(f"Error processing UI agent command: {e}", exc_info=True)
            return self._format_error_response(None, f"Error processing UI agent command: {e}")

    def _format_success_response(self, task_id: str, result: any) -> str:
        response = {
            "protocol_version": "1.0",
            "task_id": task_id,
            "responder": "UIAgent",
            "status": "success",
            "result": {
                "output": result,
                "error_message": ""
            }
        }
        return json.dumps(response, indent=2, ensure_ascii=False)

    def _format_error_response(self, task_id: str, error_message: str) -> str:
        response = {
            "protocol_version": "1.0",
            "task_id": task_id,
            "responder": "UIAgent",
            "status": "failure",
            "result": {
                "output": "",
                "error_message": error_message
            }
        }
        return json.dumps(response, indent=2, ensure_ascii=False)

if __name__ == '__main__':
    # Example usage:
    # Make sure the ui-tars-server.py is running before executing this.
    ui_agent = UIAgent()

    # Example 1: A simple task
    request_data = {
        "protocol_version": "1.0",
        "task_id": "ui-task-001",
        "requester": "RootAgent",
        "task": {
            "prompt": "Win+R, open calculator, and compute 123 + 456",
            "max_step": 50
        }
    }
    response = ui_agent.process_command(json.dumps(request_data))
    print("--- Example 1 Response ---")
    print(response)

    # Example 2: Missing prompt
    request_data_fail = {
        "protocol_version": "1.0",
        "task_id": "ui-task-002",
        "requester": "RootAgent",
        "task": {
            "max_step": 10
        }
    }
    response_fail = ui_agent.process_command(json.dumps(request_data_fail))
    print("\n--- Example 2 Response (Failure) ---")
    print(response_fail)
